import java.nio.ByteBuffer;
import java.util.Arrays;

/**
 * This class represents an IPv4 datagram.
 * 
 * @author Jonathan Jung, Timothy Hong, Ronen Adam
 *
 */
public class IPv4Datagram {

    // Maximum time-to-live value.
    public static final int INITIALIZE_TTL = 255;
    // Bitmask for do-not-fragment bit.
    private final short SET_DO_NOT_FRAGMENT = 0x4000;
    // Bitmask for set-more-bits bit.
    private final short SET_MORE_BITS = 0x2000;

    private byte[] datagramBytes;

    public IPv4Datagram(short id, boolean doNotFragment, boolean moreBits,
            short fragmentOffset, int lifetime, boolean decrementTTL,
            String sourceIp, String destinationIp, byte[] payload)
            throws InstantiationException {

        if (lifetime < 0) {
            throw new InstantiationException(
                    "IPv4 Datagram not built: Time-to-live expired.");
        } else if (decrementTTL) {
            lifetime--;

            byte timeToLive = (byte) lifetime;

            datagramBytes = assemble(id, doNotFragment, moreBits,
                    fragmentOffset, timeToLive, sourceIp, destinationIp,
                    payload);
        } else {
            byte timeToLive = (byte) lifetime;

            datagramBytes = assemble(id, doNotFragment, moreBits,
                    fragmentOffset, timeToLive, sourceIp, destinationIp,
                    payload);
        }
    }

    // TODO:
    public IPv4Datagram(byte[] data, boolean decrementTTL) {
        byte[] headerBytes = Arrays.copyOfRange(data, 0, 20);

        if (validateChecksum(headerBytes)) {
            System.out.println("a");
        }
    }

    /**
     * Assembles an IPv4 datagram. Maximum length is 65535 bytes.
     * 
     * @param id
     * @param doNotFragment
     * @param moreBits
     * @param fragmentOffset
     * @param timeToLive
     * @param sourceIp
     * @param destinationIp
     * @param payload
     * @return
     */
    private byte[] assemble(short id, boolean doNotFragment, boolean moreBits,
            short fragmentOffset, byte timeToLive, String sourceIp,
            String destinationIp, byte[] payload) {

        // The maximum length of an IPv4 datagram is 65535 bytes.
        ByteBuffer buffer = ByteBuffer.allocate(65535);

        /*
         * Append version and header length (HLEN). They are each 4 bits in
         * length so they'll need to be combined to make up 1 byte. IPv4 is the
         * first half of the byte. HLEN is the second half of the byte. IPv4 is
         * 4 hex. HLEN is 5 hex. 45 hex equals 69 decimal.
         * 
         * Byte 1. Note: array indices are (position - 1).
         */
        buffer.put(Integer.valueOf(69).byteValue());

        /*
         * Append service type, 8 bits. Value = 0 as specified in the project
         * specifications.
         * Byte 2.
         */
        buffer.put(Integer.valueOf(0).byteValue());

        /*
         * The 3rd and 4th bytes of the header are reserved for the total length
         * of the IPv4 datagram. Bytes are being added via 'relative put', so
         * just use a filler value to skip them for now.
         */
        buffer.put(Integer.valueOf(1).byteValue());
        buffer.put(Integer.valueOf(1).byteValue());

        // Append identification, 16 bits. Bytes 5 and 6.
        buffer.put(shortToBytes(id));

        /*
         * Append fragment flags and offsets. Flags = 3 bits. Fragment offset =
         * 13 bits. Combine the "fields" into a single 2-byte value.
         * Bytes 7 and 8.
         */
        int fragmentData = fragmentOffset;

        // Bitwise operation to set the appropriate bit.
        if (doNotFragment) {
            fragmentData = fragmentData | SET_DO_NOT_FRAGMENT;
        }

        // Bitwise operation to set the appropriate bit.
        if (moreBits) {
            fragmentData = fragmentData | SET_MORE_BITS;
        }
        
        buffer.putShort((short) (fragmentData));

        // Append time to live (Number of hops left), 8 bits. Byte 9.
        buffer.put(timeToLive);

        /* Append protocol, 8 bits. Value = 6, corresponding to TCP protocol.
         * Byte 10.
         */
        buffer.put(Integer.valueOf(6).byteValue());

        /*
         * The 11th and 12th bytes of the header are reserved for the checksum
         * value of the IPv4 datagram. Bytes are being added via 'relative put',
         * so just use a filler value to skip them for now.
         */
        buffer.put(Integer.valueOf(1).byteValue());
        buffer.put(Integer.valueOf(1).byteValue());

        // Append source IP, 32 bits. Bytes 13, 14, 15, and 16.
        buffer.put(ipToBytes(sourceIp));

        // Append destination IP, 32 bits. Bytes 17, 18, 19, and 20.
        buffer.put(ipToBytes(destinationIp));

        /*
         * No options or padding according to project specifications. This marks
         * the end of the header's bytes.
         */

        // Append data, variable bit length.
        buffer.put(payload);

        // Replace the filler values for 'total length' and 'checksum'.

        /*
         * Insert total length, 16 bits. Calculate the total length of the IPv4
         * datagram by using the relative position of the buffer and insert the
         * value into the buffer's bytes 2 and 3, corresponding to the 3rd and
         * 4th bytes of the header.
         */
        buffer.putShort(2, (short) buffer.position());

        // TODO: Implement the checksum method(s).
        /*
         * Insert checksum, 16 bits. Calculate the checksum of the IPv4 datagram
         * header and insert the value into the buffer's bytes 11 and 12,
         * corresponding to the 10th and 11th bytes of the header. The header
         * occupies the buffer's bytes 0 - 19.
         * 
         * Although ByteBuffer has a bulk get method available, it will alter
         * the relative put position. To prevent changes to this value, use a
         * potentially less efficient copying method.
         */
        byte[] headerBytes = new byte[20];

        for (int i = 0; i < headerBytes.length; i++) {
            headerBytes[i] = buffer.get(i);
        }

        short checksum = calculateChecksum(headerBytes);

        byte[] checksumBytes = shortToBytes(checksum);

        buffer.put(11, checksumBytes[0]);
        buffer.put(12, checksumBytes[1]);

        // Extract the datagram's bytes from the buffer.
        int size = buffer.position();

        byte[] datagramArray = new byte[size];

        for (int i = 0; i < size; i++) {
            datagramArray[i] = buffer.get(i);
        }

        return datagramArray;
    }

    /**
     * Converts a string of an IP address in do notation into an array of bytes.
     * 
     * @param ip - The IP address as a string in dot notation.
     * @return Returns an array of bytes containing the IP address.
     */
    private byte[] ipToBytes(String ip) {
        // Strip leading '/'.
        if (ip.substring(0, 1).equals("/")) {
            ip = ip.substring(1);
        }

        // Isolate each octet from the IP string.
        String[] addressSlices = ip.split("\\.");

        // Assemble an integer ip value from the octets.
        int octet1 = Integer.valueOf(addressSlices[0]);
        int octet2 = Integer.valueOf(addressSlices[1]);
        int octet3 = Integer.valueOf(addressSlices[2]);
        int octet4 = Integer.valueOf(addressSlices[3]);

        byte[] array = new byte[4];

        array[0] = Integer.valueOf(octet1).byteValue();
        array[1] = Integer.valueOf(octet2).byteValue();
        array[2] = Integer.valueOf(octet3).byteValue();
        array[3] = Integer.valueOf(octet4).byteValue();

        return array;
    }

    /**
     * Converts an array of bytes to a string of an IP address in dot notation.
     * 
     * @param data
     *            - The array containing the IP address.
     * @return Returns a string of an IP address in dot notation.
     */
    private String bytesToIP(byte[] data) {
        StringBuilder result = new StringBuilder();

        int octet = 0;

        // The first 3 octets.
        for (int i = 0; i < 3; i++) {
            octet = data[i];

            if (octet < 0) {
                octet = octet + 256;
            }

            result.append(Integer.toString(octet));
            result.append(".");
        }

        // The final octet. Separated from the first 3 to avoid a trailing '.'.
        octet = data[3];

        if (octet < 0) {
            octet = octet + 256;
        }

        result.append(Integer.toString(octet));

        result.trimToSize();

        return result.substring(0);
    }

    /**
     * Returns a string representation of the bytes in an IPv4 datagram.
     * 
     * @return Returns a string representation of an IPv4 datagram.
     */
    public String toString() {
        StringBuilder dataString = new StringBuilder();

        for (int i = 0; i < datagramBytes.length; i++) {
            dataString.append(datagramBytes[i]);
            dataString.append(" ");
        }

        dataString.trimToSize();

        return dataString.substring(0);
    }

    // TODO: calculate
    /**
     * Calculates the checksum value of an IPv4 Datagram.
     * 
     * @return
     */
    private short calculateChecksum(byte[] data) {

        /*
         * To calculate the checksum, use the procedure taken from:
         * http://en.wikipedia.org/wiki/IPv4_header_checksum
         * 
         * Step 1, calculate the sum of each 16-bit value of the header.
         */
        int sum = sumFromByteArray(data);

        /*
         * Step 2, convert the sum to binary and add the carry to the rest of
         * the sum values.
         */
        int checksum = calculateCarryPlusRestofsum(sum);

        // Step 3, flip every bit to obtain the checksum.
        checksum = ~checksum;

        // Return a 16-bit value.
        return (short) checksum;
    }

    // TODO: validate
    /**
     * Validates the checksum.
     * 
     * @param value
     * @return Returns <B>true<B> if the validation is successful.
     *         Returns <B>false<B> if the validation was not successful.
     */
    public boolean validateChecksum(byte[] data) {
        ByteBuffer buffer = ByteBuffer.allocate(data.length);

        buffer.put(data);

        int sum = sumFromByteArray(data);

        int checksum = calculateCarryPlusRestofsum(sum);

        // checksum = checksum ^ 0xFFFF0000;

        checksum = ~checksum;

        if (checksum == 0) {
            return true;
        }

        return false;
    }

    /*
     * Returns the datagram as an array of bytes.
     * 
     * @return Returns the datagram as an array of bytes.
     */
    public byte[] getBytes() {

        return datagramBytes;
    }

    // TODO: fragments
    /**
     * Fragment an IPv4Datagram into a specified array of datagrams with 
     * fragment bits and offsets calculated.
     * 
     * @param datagram - The datagram to be fragmented.
     * @param fragments - The desired number of fragments.
     * @return Returns an array of datagrams.
     */
    public IPv4Datagram[] fragment(IPv4Datagram[] datagram, int fragments) {
        IPv4Datagram[] result = {};

        return result;
    }

    /**
     * Calculates the sum of the bytes in an array of bytes.
     * 
     * @param data - The array to be summed.
     * @return Returns the sum of an array of bytes.
     */
    private int sumFromByteArray(byte[] data) {
        int[] temp = new int[data.length];

        // Correct for byte offset to avoid signed bits.
        for (int i = 0; i < data.length; i++) {
            if (data[i] < 0) {
                temp[i] = ~data[i];
            }
        }

        // Get sum of 16-bit values by assembling a hex-value pair of "bytes".
        int b1 = 0;
        int b2 = 0;
        int sum = 0;

        for (int i = 0; i < temp.length; i = i + 2) {
            b1 = temp[i];
            b2 = temp[i + 1];

            String leftByte = String.format("%2s", Integer.toHexString(b1))
                    .replace(' ', '0');
            String rightByte = String.format("%2s", Integer.toHexString(b2))
                    .replace(' ', '0');

            String hexPair = (leftByte + rightByte);

            sum += Integer.parseInt(hexPair, 16);
        }

        return sum;
    }

    /**
     * Calculates the first "nibble" and the rest of the sum for a IPv4Datagram
     * checksum.
     * 
     * @param sum - The value on which to perform the calculation.
     * @return Returns the calculated value.
     */
    private int calculateCarryPlusRestofsum(int sum) {
        /*
         * Use strings to manipulate binary values.
         * 
         * The first 4 bits are the carry and will be added to the rest of the
         * value. Compensate for the loss of leading zeros by padding zeros to
         * the left side of the string.
         */
        String sumString = String.format("%20s", Integer.toBinaryString(sum))
                .replace(' ', '0');

        // Extract the carry values.
        int carry = 0;

        String s1 = sumString.substring(0, 4);
        carry = Integer.valueOf(s1, 2);

        // Extract the rest of the sum.
        int restOfSum = 0;

        String s2 = sumString.substring(4);
        restOfSum = Integer.valueOf(s2, 2);

        // Add the carry to the rest of the sum.
        int binarySum = carry + restOfSum;

        return binarySum;
    }

    /**
     * Retrieves the bytes containing the header of an IPv4Datagram.
     * 
     * @return Returns an array containing the bytes making up the header.
     */
    private byte[] getHeader() {
        /*
         * Collect the header data.
         * 
         * Bytes ten and eleven of the IP datagram contain the checksum value
         * and must be omitted from the checksum calculation, leaving 18 bytes.
         */
        byte[] data = new byte[18];

        // Copy bytes 0 - 9.
        for (int i = 0; i < 10; i++) {
            data[i] = datagramBytes[i];
        }

        // Skip bytes 10 and 11 and continue copying bytes 12 - 19.
        for (int i = 12; i < data.length; i++) {
            data[i] = datagramBytes[i];
        }

        return data;
    }

    /**
     * Converts a 16-bit short value into an array of bytes containing two
     * 8-bit values representing the upper and lower bytes of the short.
     * 
     * @param value - The short value to be converted.
     * @return Returns and array containing the upper and lower bytes.
     */
    private byte[] shortToBytes(short value) {
        byte upper = (byte) (value >> 8);
        byte lower = (byte) (value & 0xFF);

        byte[] array = { upper, lower };

        return array;
    }
    
    /**
     * Converts an array of two bytes into a short value.
     * 
     * @param data
     * @return
     */
    private short bytesToShort(byte[] data) {
        byte upper = data[0];
        byte lower = data[1];      
        short value = 0;
        
        if (value < 0) {
            value += lower + 256;
        } else {
            value += lower;
        }
        
        value += upper << 8;
           
        return value;
    }

    /**
     * Returns the destination IP as a string in dot notation.
     * 
     * @return Returns the destination IP as a string in dot notation.
     */
    public String getDestinationIP() {
        byte[] data = Arrays.copyOfRange(datagramBytes, 16, 20);

        return bytesToIP(data);
    }

    // TODO: return values needed by router.
    /**
     * Returns the total time to live.
     * 
     * @return Returns the total time to live.
     */
    public int getTTL() {
        byte value = datagramBytes[8];

        int time = value;

        /* Due to byte overflow, the value may need to be adjusted.
         * e.g. 255 becomes -1, so (-1 + 256) yields the original value.
         */
        if (value < 0) {
            time = value + 256;
        }

        return time;
    }

    /**
     * Checks if the do-not-fragment bit of a datagram is set.
     * 
     * @return Returns <B>true</B> if the bit is set.
     *         Returns <B>false</B> if the bit is not set.
     */
    public boolean isDoNotFragmentSet() {   
        /* The do-not-fragment bit is located in the 7th byte of the header,
         * but an array starts at '0' so the bit is in byte '6'.
         */
        // Shift the bit so that the mask detects the 15th bit.
        int value = datagramBytes[6] << 8;
        
        // The do-not-fragment bit is the 15th bit.
        int fragmentBitMask = 0x00004000;
        
        // Bit-wise operation to check the 15th bit.
        value = value & fragmentBitMask;
        
        boolean isDNFSet = false;
        
        if (value == 16384) {
            isDNFSet = true;
            
        }
        
        return isDNFSet;
    }
    
    /**
     * Checks if the more-bits bit of a datagram is set.
     * 
     * @return Returns <B>true</B> if the bit is set.
     *         Returns <B>false</B> if the bit is not set.
     */
    public boolean isMoreBitSet() {   
        /* The more-fragments bit is located in the 7th byte of the header,
         * but an array starts at '0' so the bit is in byte '6'.
         */
        // Shift the bit so that the mask detects the 14th bit.
        int value = datagramBytes[6] << 8;
        
        // The do-not-fragment bit is the 14th bit.
        int fragmentBitMask = 0x00002000;
        
        // Bit-wise operation to check the 15th bit.
        value = value & fragmentBitMask;
        
        boolean isMoreBitSet = false;
        
        if (value == 8192) {
            isMoreBitSet = true;
            
        }
        
        return isMoreBitSet;
    }

    /**
     * Returns the value in the fragment offset field of the header.
     * 
     * @return Returns the value in the fragment offset field of the header.
     */
    public short getFragmentOffset() {
        // To extract the fragment offset value, turn off the flag bits.
        
        // If do-not-fragment is set, turn it off.
        int dnfBitMask = 0x00004000;
        
        // If more-bits is set, turn it off.
        int moreBitsMask = 0x00002000;

        int offset = 0;

        offset += datagramBytes[6] << 8;
//        offset += 0x000000FF;
        
        System.out.println(Integer.toHexString(offset));
        
        offset = offset | datagramBytes[7];
        
        System.out.println(Integer.toHexString(offset));
        
        offset = offset ^ dnfBitMask;
        offset = offset ^ moreBitsMask;
        
        return (short)offset;
    }

    /**
     * Returns the total length of the datagram.
     * 
     * @return Returns the total length of the datagram.
     */
    public int getLength() {

        return datagramBytes.length;
    }
}
